#include "Constants.h"
#include "RayMarch.h"
#include "MultiBlock2D.h"
#include "CudaHelpers.h"
#include "DataLoader.h"

#define _USE_MATH_DEFINES
#include <math.h>

__constant__ int d_imageWidth;
__constant__ int d_imageHeight;
__constant__ int d_imageWidthXHeight;
__constant__ int d_dataSizeX;
__constant__ int d_dataSizeY;
__constant__ int d_dataSizeZ;
__constant__ float3 d_direction;
__constant__ float3 d_up;
__constant__ float3 d_right;
__constant__ float d_scaleFactor;
__constant__ float d_eyeDistance;
__constant__ float d_thicknessScaleFactor;
__constant__ float d_inverseScaleFactor;
__constant__ float3 d_planeCenter;
__constant__ int d_raySteps;


__constant__ float d_planeDistance;

float3 cross(const float3& a, const float3& b)
{
	return make_float3(a.y * b.z - a.z * b.y, a.z * b.x - a.x * b.z, a.x * b.y - a.y * b.x);
}


void init(int imageWidth, int imageHeight, int dataSizeX, int dataSizeY, int dataSizeZ, float3 direction, float3 up, float scaleFactor, float eyeDistance, float thicknessScaleFactor)
{
	cudaMemcpyToSymbol(d_imageWidth, &imageWidth, 1 * sizeof(int));
	cudaMemcpyToSymbol(d_imageHeight, &imageHeight, 1 * sizeof(int));
	int imageWidthXHeight = imageWidth * imageHeight;
	cudaMemcpyToSymbol(d_imageWidthXHeight, &imageWidthXHeight, 1 * sizeof(int));
	cudaMemcpyToSymbol(d_dataSizeX, &dataSizeX, 1 * sizeof(int));
	cudaMemcpyToSymbol(d_dataSizeY, &dataSizeY, 1 * sizeof(int));
	cudaMemcpyToSymbol(d_dataSizeZ, &dataSizeZ, 1 * sizeof(int));
	cudaMemcpyToSymbol(d_direction, &direction, 1 * sizeof(float3));
	cudaMemcpyToSymbol(d_up, &up, 1 * sizeof(float3));
	float3 right = cross(direction, up);
	cudaMemcpyToSymbol(d_right, &right, 1 * sizeof(float3));
	cudaMemcpyToSymbol(d_scaleFactor, &scaleFactor, 1 * sizeof(float));
	float inverseScaleFactor = 1.f / scaleFactor;
	cudaMemcpyToSymbol(d_inverseScaleFactor, &inverseScaleFactor, 1 * sizeof(float));
	cudaMemcpyToSymbol(d_eyeDistance, &eyeDistance, 1 * sizeof(float));
	cudaMemcpyToSymbol(d_thicknessScaleFactor, &thicknessScaleFactor, 1 * sizeof(float));
	float planeDistance = std::max(std::max(dataSizeX, dataSizeY), dataSizeZ);
	cudaMemcpyToSymbol(d_planeDistance, &planeDistance, 1 * sizeof(float));

	float tmp = (1.414213562373095f * (planeDistance / 2.f) * scaleFactor);
	float3 planeCenter = make_float3(-direction.x * tmp, -direction.y * tmp, -direction.z * tmp);
	cudaMemcpyToSymbol(d_planeCenter, &planeCenter, 1 * sizeof(float3));

	int raySteps = (int)(planeDistance * scaleFactor);
	cudaMemcpyToSymbol(d_raySteps, &raySteps, 1 * sizeof(int));
}

__device__ float3 operator+(const float3& a, const float3& b)
{
	return make_float3(a.x + b.x, a.y + b.y, a.z + b.z);
}

__device__ void operator+=(float3& a, float3 b)
{
	a.x += b.x;
	a.y += b.y;
	a.z += b.z;
}

__device__ float3 operator*(const float3& v, float i)
{
	return make_float3(v.x * i, v.y * i, v.z * i);
}

__device__ float3 operator-(const float3& a)
{
	return make_float3(-a.x, -a.y, -a.z);
}

__device__ void operator+=(float3& a, float b)
{
	a.x += b;
	a.y += b;
	a.z += b;
}


__device__ float3 operator-(const float3& a, const float3& b)
{
	return make_float3(a.x - b.x, a.y - b.y, a.z - b.z);
}

__device__ float3 operator+(const float3& v, float i)
{
	return make_float3(v.x + i, v.y + i, v.z + i);
}

__device__ float3 operator/(const float3& v, float i)
{
	return make_float3(v.x / i, v.y / i, v.z / i);
}

__device__ void norm(float3& d)
{
	float l = sqrt(d.x * d.x + d.y * d.y + d.z * d.z);

	d.x /= l;
	d.y /= l;
	d.z /= l;
}

__global__ void rayMarch(float* data, float* output, float* tf, int tf_size, int strideX, int strideY)
{
	int xd = threadIdx.x + strideX;
	int yd = threadIdx.y + strideY;
	int x = xd - (d_imageWidth >> 1);
	int y = yd - (d_imageHeight >> 1);

	float3 plane = make_float3(d_planeCenter.x, d_planeCenter.y, d_planeCenter.z);

	plane.x += d_right.x * x + d_up.x * y;
	plane.y += d_right.y * x + d_up.y * y;
	plane.z += d_right.z * x + d_up.z * y;

	float3 d = plane - (-d_direction * d_eyeDistance);

	norm(d);

	// transpose ray to match data-cube
	plane.x += (d_dataSizeX / 2.f) * d_scaleFactor;
	plane.y += (d_dataSizeY / 2.f) * d_scaleFactor;
	plane.z += (d_dataSizeZ / 2.f) * d_scaleFactor;

	for(int i = 0;i < d_raySteps;i++)
	{
		// move ray
		plane.x += d.x;
		plane.y += d.y;
		plane.z += d.z;

		// continue if ray doesn't hit data-cube
		int dx = (int)(plane.x * d_inverseScaleFactor);
		int dy = (int)(plane.y * d_inverseScaleFactor);
		int dz = (int)(plane.z * d_inverseScaleFactor);
		if (!M_range_check(dx, d_dataSizeX) || !M_range_check(dy, d_dataSizeY) || !M_range_check(dz, d_dataSizeZ))
			continue;

		// sum pixel color
		float d = data[M_mul(dz, M_mul(d_dataSizeX, d_dataSizeY)) + M_mul(dy, d_dataSizeX) + dx];
		int tfi = ((int)(d * (tf_size >> 2)) << 2);
		int partialIndex = xd + M_mul(yd, d_imageWidth);
		float a = (1 - output[M_mul(3, d_imageWidthXHeight) + partialIndex]);
		float weight = a * // alpha
						d * // data
						tf[tfi + 3] * // transfer function
						d_thicknessScaleFactor; // scale factor
		output[partialIndex] += (1.f - tf[tfi + 0]) * weight;
		output[d_imageWidthXHeight + partialIndex] += (1.f - tf[tfi + 1]) * weight;
		output[(d_imageWidthXHeight << 1) + partialIndex] += (1.f - tf[tfi + 2]) * weight;
		output[M_mul(3, d_imageWidthXHeight) + partialIndex] += weight;

		// early ray termination
		if (a <= .0001f)
			break;
	}
}

void normalize(float3& d)
{
	float l = sqrt(d.x * d.x + d.y * d.y + d.z * d.z);

	d.x /= l;
	d.y /= l;
	d.z /= l;
}

void rotateAround(float3& a, const float3& b, float amount)
{
	float nx = b.x * (b.x * a.x + b.y * a.y + b.z * a.z) + (a.x * (b.y * b.y + b.z * b.z) - b.x * (b.y * a.y + b.z * a.z)) * cos(amount) + (-b.z * a.y + b.y * a.z) * sin(amount);
	float ny = b.y * (b.x * a.x + b.y * a.y + b.z * a.z) + (a.y * (b.x * b.x + b.z * b.z) - b.y * (b.x * a.x + b.z * a.z)) * cos(amount) + (b.z * a.x - b.x * a.z) * sin(amount);
	float nz = b.z * (b.x * a.x + b.y * a.y + b.z * a.z) + (a.z * (b.x * b.x + b.y * b.y) - b.z * (b.x * a.x + b.y * a.y)) * cos(amount) + (-b.y * a.x + b.x * a.y) * sin(amount);

	a.x = nx;
	a.y = ny;
	a.z = nz;
}

void CheckKernelError(const string& kernelName)
{
	cudaError_t i = cudaGetLastError();
	if (i)
		printf(("Kernel (\"" + kernelName + "\") error: %s\n").c_str(), cudaGetErrorString(i));
}

void rayMarch(float eyeDistance, const float3& direction, float* data, float* output, Settings* settings)
{
	// calculate up vector
	float3 right = make_float3(direction.x, .0f, direction.z);
	normalize(right);
	rotateAround(right, make_float3(.0f, 1.f, .0f), M_PI / 2.f);
	float3 newUp = make_float3(direction.y * right.z - right.y * direction.z, direction.z * right.x - right.z * direction.x, direction.x * right.y - right.x * direction.y);

	init(settings->getImageWidth(), settings->getImageHeight(), settings->getDataSizeX(), settings->getDataSizeY(), settings->getDataSizeZ(), direction, newUp, settings->getScaleFactor(), eyeDistance, settings->getDivisor());

	// TF
	float* tf;
	int tf_size;
	LoadTF(settings->getTFFilename(), settings->getTFWeights(), &tf, &tf_size);

	MultiBlock2D mb(settings->getImageWidth(), settings->getImageHeight());
	for(int i = 0;i < mb.numX;i++)
		for(int j = 0;j < mb.numY;j++)
		{
			dim3 dimBlock = mb.GetBlockSize(i, j);
			dim3 dimGrid;

			//cudaPrintfInit();
			rayMarch<<<dimGrid, dimBlock>>>(data, output, tf, tf_size, i * mb.StrideX(), j * mb.StrideY());
			CheckKernelError("rayMarch");
			//cudaPrintfDisplay(stdout, true);
			//cudaPrintfEnd();

		}

	Free(tf);
}
